'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: file_extractor.py
* @Time: 2025/11/12
* @All Rights Reserve By Brtc
'''
import os
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Union

import requests
from injector import inject
from langchain_community.document_loaders import UnstructuredFileLoader, UnstructuredMarkdownLoader, \
    UnstructuredPowerPointLoader, UnstructuredXMLLoader, UnstructuredCSVLoader, UnstructuredPDFLoader, \
    UnstructuredExcelLoader, UnstructuredWordDocumentLoader
from langchain_core.documents import Document as LCDocument

from internal.model import UploadFile
from internal.service import CosService


@inject
@dataclass
class FileExtractor:
    """文件提取,用于将云端文件 upload 记录加载成langchain 对应的文档加载器"""
    cos_service:CosService


    def load(self,
             upload_file:UploadFile,
             return_text:bool = False,
             is_unstructured:bool = True)->Union[list[LCDocument],str]:
        """加载传入的upload file 记录返回langchain 文档列表或者字符串"""
        #1、创建一个临时文件夹
        with tempfile.TemporaryDirectory() as temp_dir:
            #2、构建一个临时文件夹路径
            file_path = os.path.join(temp_dir,os.path.basename(upload_file.key))
            #3、将对象存储的文件下载到本地
            self.cos_service.download_file(upload_file.key, file_path)
            #4、从指定的路径中取加载文件
            return self.load_from_file(file_path, return_text, is_unstructured)


    @classmethod
    def load_from_url(cls, url:str, return_text:bool = False)->Union[list[LCDocument],str]:
        """从传入的URL中取加载数据, 返回langchain 的的文档列表或者字符串"""
        #1、下载远程的URL文件
        response = requests.get(url)
        with tempfile.TemporaryDirectory() as temp_dir:
            # 2、构建一个临时文件夹路径
            file_path = os.path.join(temp_dir, os.path.basename(url))
            # 3、将对象存储的文件下载到本地
            with open(file_path,'wb') as f:
                f.write(response.content)
            # 4、从指定的路径中取加载文件
            return cls.load_from_file(file_path, return_text)


    @classmethod
    def load_from_file(cls, file_path:str, return_text:bool = False, is_unstructured:bool = True)->Union[list[LCDocument],str]:
        """从本地文件中加载数据返回langchain的文档列表或者字符串"""
        #1、获取文件扩展名
        delimiter = "\n\n"
        file_extension = Path(file_path).suffix.lower()

        #2、根据不同 文件扩展名加载不同的加载器
        if file_extension in [".doc",".docx"]:
            loader = UnstructuredWordDocumentLoader(file_path)
        elif file_extension in [".xls",".xlsx"]:
            loader = UnstructuredExcelLoader(file_path)
        elif file_extension == ".pdf" :
            loader = UnstructuredPDFLoader(file_path)
        elif file_extension == ".csv":
            loader = UnstructuredCSVLoader(file_path)
        elif file_extension == ".xml":
            loader = UnstructuredXMLLoader(file_path)
        elif file_extension in [".ppt",".pptx"]:
            loader = UnstructuredPowerPointLoader(file_path)
        elif file_extension in [".md",".markdown"]:
            loader = UnstructuredMarkdownLoader(file_path)
        else:
            loader = UnstructuredFileLoader(file_path)
        #3、返回加载的文档列表
        return delimiter.join([document.page_content for document in loader.load()]) if return_text else loader.load()
